package edge

import (
	"context"
	"fmt"

	"github.com/DataDog/KubeHound/pkg/kubehound/graph/adapter"
	"github.com/DataDog/KubeHound/pkg/kubehound/graph/types"
	"github.com/DataDog/KubeHound/pkg/kubehound/models/converter"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/cache"
	"github.com/DataDog/KubeHound/pkg/kubehound/storage/storedb"
	"github.com/DataDog/KubeHound/pkg/kubehound/store/collections"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

func init() {
	Register(&EndpointExploitExternal{}, RegisterDefault)
}

type EndpointExploitExternal struct {
	BaseEdge
}

type sliceEndpointGroup struct {
	Endpoint  primitive.ObjectID `bson:"_id" json:"endpoint_id"`
	Container primitive.ObjectID `bson:"container_id" json:"container_id"`
}

func (e *EndpointExploitExternal) Label() string {
	return "ENDPOINT_EXPLOIT"
}

func (e *EndpointExploitExternal) Name() string {
	return "EndpointExploitExternal"
}

func (e *EndpointExploitExternal) AttckTechniqueID() AttckTechniqueID {
	return AttckTechniqueExploitationOfRemoteServices
}

func (e *EndpointExploitExternal) AttckTacticID() AttckTacticID {
	return AttckTacticLateralMovement
}

func (e *EndpointExploitExternal) Processor(ctx context.Context, oic *converter.ObjectIDConverter, entry any) (any, error) {
	typed, ok := entry.(*sliceEndpointGroup)
	if !ok {
		return nil, fmt.Errorf("invalid type passed to processor: %T", entry)
	}

	return adapter.GremlinEdgeProcessor(ctx, oic, e.Label(), typed.Endpoint, typed.Container, map[string]any{
		"attckTechniqueID": string(e.AttckTechniqueID()),
		"attckTacticID":    string(e.AttckTacticID()),
	})
}

func (e *EndpointExploitExternal) Stream(ctx context.Context, store storedb.Provider, c cache.CacheReader,
	callback types.ProcessEntryCallback, complete types.CompleteQueryCallback) error {

	endpoints := adapter.MongoDB(ctx, store).Collection(collections.EndpointName)

	// K8s endpoint slices must be ingested before containers. In this stage we need to match store.Endpoint documents that
	// are generated via K8s EndpointSlice objects and match them to the container exposing the endpoint. The other case of
	// store.Endpoint documents not associated with an EndpointSlice is handled separately.
	pipeline := []bson.M{
		{
			// Match only endpoints with a matching EndpointSlice
			"$match": bson.M{
				"has_slice":            true,
				"runtime.runID":        e.runtime.RunID.String(),
				"runtime.cluster.name": e.runtime.Cluster.Name,
			},
		},
		{
			// Lookup the container matching the slice. This requires a match on namespace/pod/port/protocol
			"$lookup": bson.M{
				"as":   "matchContainers",
				"from": "containers",
				"let": bson.M{
					"pod":   "$pod_name",
					"podNS": "$pod_namespace",
					"port":  "$port.port",
					"proto": "$port.protocol",
				},
				"pipeline": []bson.M{
					{
						"$match": bson.M{
							"$expr": bson.M{
								"$and": bson.A{
									bson.M{"$eq": bson.A{
										"$inherited.namespace", "$$podNS",
									}},
									bson.M{"$eq": bson.A{
										"$inherited.pod_name", "$$pod",
									}},
									bson.M{"$ne": bson.A{
										"$k8.ports", nil,
									}},
									// Cannot use an $elemMatch with pipeline variables so use the more convoluted $filter syntax to match container port/protocol
									// See: https://www.mongodb.com/community/forums/t/equivalent-of-elemmatch-query-operator-for-use-in-match-within-the-aggregation-lookup-with-pipeline/5360
									bson.M{"$gt": bson.A{
										bson.M{"$size": bson.M{"$filter": bson.M{
											"input": "$k8.ports",
											"as":    "p",
											"cond": bson.M{
												"$and": bson.A{
													bson.M{"$eq": bson.A{
														"$$p.containerport", "$$port",
													}},
													bson.M{"$eq": bson.A{
														"$$p.protocol", "$$proto",
													}},
												}},
										}}},
										0,
									}},
								},
							},
							"runtime.runID":        e.runtime.RunID.String(),
							"runtime.cluster.name": e.runtime.Cluster.Name,
						},
					},
					{
						"$project": bson.M{
							"_id": 1,
						},
					},
				},
			},
		},
		{
			"$unwind": "$matchContainers",
		},
		{
			"$project": bson.M{
				"_id":          1,
				"container_id": "$matchContainers._id",
			},
		},
	}

	cur, err := endpoints.Aggregate(ctx, pipeline)
	if err != nil {
		return err
	}
	defer cur.Close(ctx)

	return adapter.MongoCursorHandler[sliceEndpointGroup](ctx, cur, callback, complete)
}
